package com.ecar.oauth.provider.beans;

import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.common.DefaultExpiringOAuth2RefreshToken;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.DefaultOAuth2RefreshToken;
import org.springframework.security.oauth2.common.ExpiringOAuth2RefreshToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.common.exceptions.InvalidGrantException;
import org.springframework.security.oauth2.common.exceptions.InvalidScopeException;
import org.springframework.security.oauth2.common.exceptions.InvalidTokenException;
import org.springframework.security.oauth2.common.exceptions.OAuth2Exception;
import org.springframework.security.oauth2.provider.ClientDetails;
import org.springframework.security.oauth2.provider.ClientDetailsService;
import org.springframework.security.oauth2.provider.ClientRegistrationException;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.OAuth2Request;
import org.springframework.security.oauth2.provider.TokenRequest;
import org.springframework.security.oauth2.provider.token.DefaultTokenServices;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.security.web.authentication.preauth.PreAuthenticatedAuthenticationToken;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.Assert;

import java.util.Date;
import java.util.Set;
import java.util.UUID;

/**
 * Created by 30 on 2016/12/22.
 */
public class OAuthTokenServices extends DefaultTokenServices {

	private int refreshTokenValiditySeconds = 2592000;
	private int accessTokenValiditySeconds = 'ꣀ';
	private boolean supportRefreshToken = false;
	private boolean reuseRefreshToken = true;
	private TokenStore tokenStore;
	private ClientDetailsService clientDetailsService;
	private TokenEnhancer accessTokenEnhancer;
	private AuthenticationManager authenticationManager;

	public OAuthTokenServices() {
	}

	public void afterPropertiesSet() throws Exception {
		Assert.notNull(tokenStore, "tokenStore must be set");
	}

	@Transactional
	public OAuth2AccessToken createAccessToken(OAuth2Authentication authentication) {
		OAuth2AccessToken existingAccessToken = tokenStore.getAccessToken(authentication);
		OAuth2RefreshToken refreshToken = null;
		if (existingAccessToken != null) {
			if (!existingAccessToken.isExpired()) {
				tokenStore.storeAccessToken(existingAccessToken, authentication);
				return existingAccessToken;
			}

			if (existingAccessToken.getRefreshToken() != null) {
				refreshToken = existingAccessToken.getRefreshToken();
				tokenStore.removeRefreshToken(refreshToken);
			}

			tokenStore.removeAccessToken(existingAccessToken);
		}

		if (refreshToken == null) {
			refreshToken = createRefreshToken(authentication);
		} else if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
			ExpiringOAuth2RefreshToken accessToken = (ExpiringOAuth2RefreshToken) refreshToken;
			if (System.currentTimeMillis() > accessToken.getExpiration().getTime()) {
				refreshToken = createRefreshToken(authentication);
			}
		}

		OAuth2AccessToken accessToken = createAccessToken(authentication, refreshToken);
		tokenStore.storeAccessToken(accessToken, authentication);
		refreshToken = accessToken.getRefreshToken();
		if (refreshToken != null) tokenStore.storeRefreshToken(refreshToken, authentication);

		return accessToken;
	}

	@Transactional(noRollbackFor = {InvalidTokenException.class, InvalidGrantException.class})
	public OAuth2AccessToken refreshAccessToken(String refreshTokenValue, TokenRequest tokenRequest) throws OAuth2Exception {
		if (!supportRefreshToken) throw new InvalidGrantException("Invalid refresh token: " + refreshTokenValue);

		OAuth2RefreshToken refreshToken = tokenStore.readRefreshToken(refreshTokenValue);
		if (refreshToken == null) throw new InvalidGrantException("Invalid refresh token: " + refreshTokenValue);

		OAuth2Authentication authentication = tokenStore.readAuthenticationForRefreshToken(refreshToken);
		if (authenticationManager != null && !authentication.isClientOnly()) {
			PreAuthenticatedAuthenticationToken clientId =
				new PreAuthenticatedAuthenticationToken(authentication.getUserAuthentication(), "",
														authentication.getAuthorities());
			Authentication clientId1 = authenticationManager.authenticate(clientId);
			Object accessToken = authentication.getDetails();
			authentication = new OAuth2Authentication(authentication.getOAuth2Request(), clientId1);
			authentication.setDetails(accessToken);
		}

		String clientId2 = authentication.getOAuth2Request().getClientId();
		if (clientId2 != null && clientId2.equals(tokenRequest.getClientId())) {
			tokenStore.removeAccessTokenUsingRefreshToken(refreshToken);
			if (isExpired(refreshToken)) {
				tokenStore.removeRefreshToken(refreshToken);
				throw new InvalidTokenException("Invalid refresh token (expired): " + refreshToken);
			}

			authentication = createRefreshedAuthentication(authentication, tokenRequest);
			if (!reuseRefreshToken) {
				tokenStore.removeRefreshToken(refreshToken);
				refreshToken = createRefreshToken(authentication);
			}

			OAuth2AccessToken accessToken1 = createAccessToken(authentication, refreshToken);
			tokenStore.storeAccessToken(accessToken1, authentication);
			if (!reuseRefreshToken) {
				tokenStore.storeRefreshToken(accessToken1.getRefreshToken(), authentication);
			}

			return accessToken1;
		}

		throw new InvalidGrantException("Wrong client for this refresh token: " + refreshTokenValue);
	}

	public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication) {
		return tokenStore.getAccessToken(authentication);
	}

	private OAuth2Authentication createRefreshedAuthentication(OAuth2Authentication authentication, TokenRequest request)
		throws OAuth2Exception {
		Set scope = request.getScope();
		OAuth2Request clientAuth = authentication.getOAuth2Request().refresh(request);
		if (scope != null && !scope.isEmpty()) {
			Set originalScope = clientAuth.getScope();
			if (originalScope == null || !originalScope.containsAll(scope)) {
				throw new InvalidScopeException("Unable to narrow the scope of the client authentication to " + scope + ".",
												originalScope);
			}

			clientAuth = clientAuth.narrowScope(scope);
		}

		return new OAuth2Authentication(clientAuth, authentication.getUserAuthentication());
	}

	protected boolean isExpired(OAuth2RefreshToken refreshToken) {
		if (!(refreshToken instanceof ExpiringOAuth2RefreshToken)) {
			return false;
		} else {
			ExpiringOAuth2RefreshToken expiringToken = (ExpiringOAuth2RefreshToken) refreshToken;
			return expiringToken.getExpiration() == null || System.currentTimeMillis() > expiringToken.getExpiration().getTime();
		}
	}

	public OAuth2AccessToken readAccessToken(String accessToken) {
		return tokenStore.readAccessToken(accessToken);
	}

	public OAuth2Authentication loadAuthentication(String accessTokenValue) throws InvalidTokenException {
		OAuth2AccessToken accessToken = this.tokenStore.readAccessToken(accessTokenValue);
		if (accessToken == null) {
			throw new InvalidTokenException("Invalid access token: " + accessTokenValue);
		} else if (accessToken.isExpired()) {
			tokenStore.removeAccessToken(accessToken);
			throw new InvalidTokenException("Access token expired: " + accessTokenValue);
		} else {
			OAuth2Authentication result = tokenStore.readAuthentication(accessToken);
			if (result == null) {
				throw new InvalidTokenException("Invalid access token: " + accessTokenValue);
			} else {
				if (clientDetailsService != null) {
					String clientId = result.getOAuth2Request().getClientId();

					try {
						clientDetailsService.loadClientByClientId(clientId);
					} catch (ClientRegistrationException var6) {
						throw new InvalidTokenException("Client not valid: " + clientId, var6);
					}
				}

				return result;
			}
		}
	}

	public String getClientId(String tokenValue) {
		OAuth2Authentication authentication = this.tokenStore.readAuthentication(tokenValue);
		if (authentication == null) {
			throw new InvalidTokenException("Invalid access token: " + tokenValue);
		} else {
			OAuth2Request clientAuth = authentication.getOAuth2Request();
			if (clientAuth == null) {
				throw new InvalidTokenException("Invalid access token (no client id): " + tokenValue);
			} else {
				return clientAuth.getClientId();
			}
		}
	}

	public boolean revokeToken(String tokenValue) {
		OAuth2AccessToken accessToken = this.tokenStore.readAccessToken(tokenValue);
		if (accessToken == null) {
			return false;
		} else {
			if (accessToken.getRefreshToken() != null) {
				tokenStore.removeRefreshToken(accessToken.getRefreshToken());
			}

			tokenStore.removeAccessToken(accessToken);
			return true;
		}
	}

	private OAuth2RefreshToken createRefreshToken(OAuth2Authentication authentication) {
		if (!isSupportRefreshToken(authentication.getOAuth2Request())) {
			return null;
		} else {
			int validitySeconds = getRefreshTokenValiditySeconds(authentication.getOAuth2Request());
			String value = UUID.randomUUID().toString();
			return validitySeconds > 0 ? new DefaultExpiringOAuth2RefreshToken(value, new Date(
				System.currentTimeMillis() + (long) validitySeconds * 1000L)) : new DefaultOAuth2RefreshToken(value);
		}
	}

	private OAuth2AccessToken createAccessToken(OAuth2Authentication authentication, OAuth2RefreshToken refreshToken) {
		DefaultOAuth2AccessToken token = new DefaultOAuth2AccessToken(UUID.randomUUID().toString());
		int validitySeconds = getAccessTokenValiditySeconds(authentication.getOAuth2Request());
		if (validitySeconds > 0) {
			token.setExpiration(new Date(System.currentTimeMillis() + (long) validitySeconds * 1000L));
		}

		token.setRefreshToken(refreshToken);
		token.setScope(authentication.getOAuth2Request().getScope());
		return accessTokenEnhancer != null ? accessTokenEnhancer.enhance(token, authentication) : token;
	}

	protected int getAccessTokenValiditySeconds(OAuth2Request clientAuth) {
		if (this.clientDetailsService != null) {
			ClientDetails client = clientDetailsService.loadClientByClientId(clientAuth.getClientId());
			Integer validity = client.getAccessTokenValiditySeconds();
			if (validity != null) {
				return validity;
			}
		}

		return accessTokenValiditySeconds;
	}

	protected int getRefreshTokenValiditySeconds(OAuth2Request clientAuth) {
		if (clientDetailsService != null) {
			ClientDetails client = clientDetailsService.loadClientByClientId(clientAuth.getClientId());
			Integer validity = client.getRefreshTokenValiditySeconds();
			if (validity != null) {
				return validity;
			}
		}

		return refreshTokenValiditySeconds;
	}

	protected boolean isSupportRefreshToken(OAuth2Request clientAuth) {
		if (clientDetailsService != null) {
			ClientDetails client = clientDetailsService.loadClientByClientId(clientAuth.getClientId());
			return client.getAuthorizedGrantTypes().contains("refresh_token");
		} else {
			return supportRefreshToken;
		}
	}

	public void setTokenEnhancer(TokenEnhancer accessTokenEnhancer) {
		this.accessTokenEnhancer = accessTokenEnhancer;
	}

	public void setRefreshTokenValiditySeconds(int refreshTokenValiditySeconds) {
		this.refreshTokenValiditySeconds = refreshTokenValiditySeconds;
	}

	public void setAccessTokenValiditySeconds(int accessTokenValiditySeconds) {
		this.accessTokenValiditySeconds = accessTokenValiditySeconds;
	}

	public void setSupportRefreshToken(boolean supportRefreshToken) {
		this.supportRefreshToken = supportRefreshToken;
	}

	public void setReuseRefreshToken(boolean reuseRefreshToken) {
		this.reuseRefreshToken = reuseRefreshToken;
	}

	public void setTokenStore(TokenStore tokenStore) {
		this.tokenStore = tokenStore;
	}

	public void setAuthenticationManager(AuthenticationManager authenticationManager) {
		this.authenticationManager = authenticationManager;
	}

	public void setClientDetailsService(ClientDetailsService clientDetailsService) {
		this.clientDetailsService = clientDetailsService;
	}
}
